import numpy as np
from .module import Module
from .utils import im2col


class Linear(Module):
    def __init__(self, in_features, out_features):
        super(Linear, self).__init__()
        self.__w = np.random.randn(out_features, in_features)
        self.__b = np.zeros(out_features)

    def forward(self, x):
        x = np.dot(self.__w, x) + self.__b
        return x

    def get_weights(self):
        return [self.__w, self.__b]

    def set_weights(self, weight):
        assert self.__w.shape == weight[0].shape, 'Size mismatch.'
        assert self.__b.shape == weight[1].shape, 'Size mismatch.'
        self.__w = weight[0]
        self.__b = weight[1]
        return 2


class Conv2d(Module):
    def __init__(self, in_channels, out_channels, kernel_size=(3, 3), padding=0):
        super(Conv2d, self).__init__()
        self.__w = np.random.randn(out_channels, in_channels, kernel_size[0], kernel_size[1])
        self.__b = np.zeros(out_channels)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.padding = padding

    def forward(self, x):
        xc, xh, xw = x.shape
        kh, kw = self.kernel_size
        h = xh + self.padding * 2 - kh + 1
        w = xw + self.padding * 2 - kw + 1
        if self.padding!=0:
            y = np.zeros([xc, xh + self.padding * 2, xw + self.padding * 2])
            y[:, 1:-1, 1:-1] = x
        else:
            y = x
        xflat = im2col(y, self.kernel_size)
        y = np.zeros([self.out_channels, h, w])
        for l in range(self.out_channels):
            kernel = self.__w[l, :, :, :].flatten()
            y[l] = np.matmul(xflat, kernel).reshape(xh, xw)
        return y

    def get_weights(self):
        return [self.__w, self.__b]

    def set_weights(self, weight):
        assert self.__w.shape == weight[0].shape, 'Size mismatch.'
        assert self.__b.shape == weight[1].shape, 'Size mismatch.'
        self.__w = weight[0]
        self.__b = weight[1]
        return 2


class LeakyReLU(Module):
    def __init__(self, negative_slope=0.01):
        super(LeakyReLU, self).__init__()
        self.negative_slope = negative_slope

    def forward(self, x):
        mask = x < 0
        x[mask] *= self.negative_slope
        return x


class MaxPool2d(Module):
    def __init__(self, kernel_size, stride):
        super(MaxPool2d, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride

    def forward(self, x):
        y = np.zeros([x.shape[0], x.shape[1] // 2, x.shape[2] // 2])
        for n in range(x.shape[0]):
            xreshape = np.stack([x[n, ::2, ::2], x[n, 1::2, ::2], x[n, ::2, 1::2], x[n, 1::2, 1::2]], axis=0)
            y[n, :, :] = np.max(xreshape, axis=0)
        return y
